import time


import cv2
import numpy as np
import pymysql


def load_md(graph_txt, inference_pb):
    inference_pb = "D:/diansai/test/frozen_inference_graph.pb"  # pb
    graph_txt = "D:/diansai/test/graph.pbtxt"                   # txt
    net = cv2.dnn.readNetFromTensorflow(inference_pb, graph_txt)    # 旧版 需要手动在cv22文件夹里更换cv22.cp37-win_amd64.pyd      10fps
    # net = cv2.dnn.readNetFromModelOptimizer(graph_xml, inference_bin)  # openvino 需要手动在cv22文件夹里更换cv22.pyd  60fps
    # net.setPreferableBackend(DNN_BACKEND_INFERENCE_ENGINE)  # openvino加速    # 默认设置 可以不添加
    # net.setPreferableTarget(DNN_TARGET_CPU)  # 。。
    return net


def Video_create(url):
    window_name = url
    if url == '0':
        url = 0
    elif url == '1':
        url = 1

    cap = cv2.VideoCapture(url)
    fps = cap.get(5)
    print(fps)
    if fps == 0.0:  # fps的 可能处理
        fps = 25
    window_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) / 10) * 20, int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) / 10) * 20)
    i = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            print('读不到有效帧')
            cap.release()
            cv2.destroyAllWindows()
            break
        frame = cv2.resize(frame, window_size)
        start_time = time.time()
        model_return(net, frame)
        fps = 1 / float(time.time() - start_time)
        # print(fps)
        cv2.putText(frame, f"fps:{fps:4.2f}", (20, 40), 1, 1.3, (0, 0, 255), 1)
        cv2.imshow(window_name, frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):  # 一般视频 为 25fps  int(1000 / fps)
            cap.release()
            cv2.destroyAllWindows()
        i += 1


def model_return(net, frame):
    h, w = frame.shape[:2]
    im_tensor = cv2.dnn.blobFromImage(frame, size=(300, 300), swapRB=True, crop=False)
    net.setInput(im_tensor)
    cvOut = net.forward()

    for detect in cvOut[0, 0, :, :]:
        # detect = detect.tolist()  # numpy  ->  list
        score = detect[2]
        left = detect[3] * w
        top = detect[4] * h
        right = detect[5] * w
        bottom = detect[6] * h

        if int(detect[1]) == 1:  # 画框
            if score > 0.9:
                cv2.putText(frame, str(1),(int(left), int(top)), 1, 3, (0, 255, 255), 1)
                cv2.rectangle(frame, (int(left), int(top)), (int(right), int(bottom)), (0, 255, 255), 4)

        elif int(detect[1]) == 2:
            if score > 0.8:
                cv2.putText(frame, str(2), (int(left), int(top)), 1, 1.2, (255, 0, 255), 2)
                cv2.rectangle(frame, (int(left), int(top)), (int(right), int(bottom)), (255, 0, 255), 4)

                # cv2.rectangle(im_tensor, rect, Scalar(0, 0, 255), 2, 8, 0)
                # putText(image, format(" baomihua %.2f, %s", confidence, "face"), Point(left - 10, top - 5),FONT_HERSHEY_SIMPLEX, 0.7, Scalar(0, 0, 255), 2, 8)
                # check_data(detect)
                # continue
        # elif int(detect[1]) == 2:
        #     if score > 0.91:
        #         check_data(detect)
        #         continue
        # elif int(detect[1]) == 3:
        #     if score > 0.91:
        #         check_data(detect)
        #         continue


def check_data(data):
    # print(data)
    pass


if __name__ == '__main__':
    inference_pb = "D:/diansai/test/frozen_inference_graph.pb"  # pb
    graph_txt = "D:/diansai/test/graph.pbtxt"                   # pbtxt
    net = load_md(graph_txt, inference_pb)
    url = '0'
    Video_create(url)